#include "lane_seg.h"
#include <chrono>
#include <fstream>
#include <opencv2/opencv.hpp>

#define LANE_NUM_CLASSES (5)
using namespace perception::camera;

std::shared_ptr<LaneSeg> LaneSeg::detector_ = nullptr;

LaneSeg::LaneSeg()
    : logger__(rclcpp::get_logger("LaneSeg"))
{
}

bool LaneSeg::init(const std::string& model_dir)
{
    num_class_ = LANE_NUM_CLASSES;
    const std::string engine_file_path = model_dir + "/lane_seg.engine";
    RCLCPP_INFO(logger__, "engine file: %s", engine_file_path.c_str());
    std::ifstream model_file(engine_file_path, std::ios::binary);
    if (!model_file.good()) {
        RCLCPP_FATAL(logger__, "model file err:%s", engine_file_path.c_str());
    }

    model_file.seekg(0, std::ios::end);
    auto size = model_file.tellg();
    model_file.seekg(0, std::ios::beg);
    char* trtModelStream = new char[size];
    assert(trtModelStream);
    model_file.read(trtModelStream, size);
    model_file.close();
    initLibNvInferPlugins(&this->g_logger_, "");
    this->runtime_ = nvinfer1::createInferRuntime(this->g_logger_);
    assert(this->runtime_ != nullptr);

    this->engine_ = this->runtime_->deserializeCudaEngine(trtModelStream, size);
    assert(this->engine_ != nullptr);
    delete[] trtModelStream;
    this->context_ = this->engine_->createExecutionContext();

    assert(this->context_ != nullptr);
    cudaStreamCreate(&this->stream_);
    this->num_bindings = this->engine_->getNbBindings();
    assert(this->num_bindings == 2);
    for (int i = 0; i < this->num_bindings; ++i) {
        Binding binding;
        nvinfer1::Dims dims;
        nvinfer1::DataType dtype = this->engine_->getBindingDataType(i);
        std::string name = this->engine_->getBindingName(i);
        binding.name = name;
        binding.dsize = type_to_size(dtype);

        bool IsInput = engine_->bindingIsInput(i);
        if (IsInput) {
            this->num_inputs += 1;
            dims = this->engine_->getProfileDimensions(
                i,
                0,
                nvinfer1::OptProfileSelector::kMAX);
            std::cout << i << ", getProfileDimensions:" << dims.d[0] << "," << dims.d[1] << "," << dims.d[2] << "," << dims.d[3] << std::endl;
            std::cout << "ndims:" << dims.nbDims << std::endl;
            // dims.d[2] = infer_h_;
            // dims.d[3] = infer_w_;
            binding.size = get_size_by_dims(dims);
            binding.dims = dims;
            this->input_bindings.push_back(binding);
            // set max opt shape
            this->context_->setBindingDimensions(i, dims);
        } else {
            dims = this->context_->getBindingDimensions(i);
            std::cout << i << ",output getProfileDimensions:" << dims.d[0] << "," << dims.d[1] << "," << dims.d[2] << "," << dims.d[3] << std::endl;
            std::cout << "output ndims:" << dims.nbDims << std::endl;
            binding.size = get_size_by_dims(dims);
            binding.dims = dims;
            this->output_bindings.push_back(binding);
            this->num_outputs += 1;
        }
    }
    data_ = new (float[infer_w_ * infer_h_ * 3]);
    score_map_ = cv::Mat(infer_h_, infer_w_, CV_32FC(LANE_NUM_CLASSES), cv::Scalar(0, 0, 0));
    // lane_seg_ = cv::Mat(infer_h_, infer_w_, CV_8FC3, cv::Scalar(0, 0, 0));
    make_pipe(true);
    return true;
}

void LaneSeg::make_pipe(bool warmup)
{

    for (auto& bindings : this->input_bindings) {
        void* d_ptr { nullptr };
        cudaError_t en = cudaMalloc(
            &d_ptr,
            bindings.size * bindings.dsize);
        RCLCPP_INFO(logger__, "errno:%d", en);
        this->device_ptrs.push_back(d_ptr);
    }

    for (auto& bindings : this->output_bindings) {
        void *d_ptr { nullptr }, *h_ptr { nullptr };
        // size_t size = 29*8400*sizeof(float);//640 * 640 * 3 * sizeof(float);//
        size_t size = bindings.size * bindings.dsize;
        RCLCPP_INFO(logger__, "lane_seg output bindings size:%ld,dsize:%ld", bindings.size, bindings.dsize);
        cudaError_t en = cudaMalloc(
            &d_ptr,
            size);
        RCLCPP_INFO(logger__, "cudama:%d", en);
        en = cudaHostAlloc(
            &h_ptr,
            size,
            0);
        RCLCPP_INFO(logger__, "host:%d", en);
        this->device_ptrs.push_back(d_ptr);
        this->host_ptrs.push_back(h_ptr);
    }
    if (warmup) {
        for (int i = 0; i < 10; i++) {
            for (auto& bindings : this->input_bindings) {
                // size_t size = 640 * 640 * 3 * sizeof(float);//
                size_t size = bindings.size * bindings.dsize;
                void* h_ptr = malloc(size);
                memset(h_ptr, 0, size);
                // cudaError_t en =
                cudaMemcpy(
                    this->device_ptrs[0],
                    h_ptr,
                    size,
                    cudaMemcpyHostToDevice);
                free(h_ptr);
            }
            this->infer();
        }
        RCLCPP_INFO(logger__, "model warmup 10 times2");
    }
}

cv::Mat LaneSeg::resize_keep_ratio(const cv::Mat& input, const cv::Size& dstSize, const cv::Scalar& bgcolor)
{
    cv::Mat output;

    double h1 = dstSize.width * (input.rows / (double)input.cols);
    double w2 = dstSize.height * (input.cols / (double)input.rows);
    if (h1 <= dstSize.height) {
        cv::resize(input, output, cv::Size(dstSize.width, h1));
    } else {
        cv::resize(input, output, cv::Size(w2, dstSize.height));
    }

    int top = (dstSize.height - output.rows) / 2;
    int down = (dstSize.height - output.rows + 1) / 2;
    int left = (dstSize.width - output.cols) / 2;
    int right = (dstSize.width - output.cols + 1) / 2;
    pad_top_ = top;
    pad_left_ = left;
    cv::copyMakeBorder(output, output, top, down, left, right, cv::BORDER_CONSTANT, bgcolor);

    return output;
}

void LaneSeg::deinit()
{
    this->context_->destroy();
    this->engine_->destroy();
    this->runtime_->destroy();
    cudaStreamDestroy(this->stream_);
    delete[] data_;
    for (auto& ptr : this->device_ptrs) {
        assert(cudaFree(ptr));
    }

    for (auto& ptr : this->host_ptrs) {
        assert(cudaFreeHost(ptr));
    }

    detector_ = nullptr;
    return;
}

// 只支持1920x1080图片，不是这个尺寸会resize
void LaneSeg::copy_from_mat(const cv::Mat& image)
{
    // auto time_cpy_start = std::chrono::high_resolution_clock::now();
    // cv::Mat blob;
    std::vector<float> mean_value { 0.406, 0.456, 0.485 }; // BGR
    std::vector<float> std_value { 0.225, 0.224, 0.229 };

    // auto tt0 = std::chrono::system_clock::now();
    // int infer_h = 360, infer_w = 640;
    image_orig_ = image.clone();

    // cv::resize(image, blob, { infer_w_, infer_h_ }, );
    // image_resize_ = blob.clone();
    cv::Mat pr_img = resize_keep_ratio(image, cv::Size(infer_w_, infer_h_),
        cv::Scalar { 0, 0, 0 });
    int i = 0;
    for (int row = 0; row < pr_img.rows; ++row) {
        uchar* uc_pixel = pr_img.data + row * pr_img.step;
        for (int col = 0; col < pr_img.cols; ++col) {
            data_[i] = (uc_pixel[2] / 255.0 - mean_value[2]) / std_value[2];
            data_[i + pr_img.rows * pr_img.cols] = (uc_pixel[1] / 255.0 - mean_value[1]) / std_value[1];
            data_[i + 2 * pr_img.rows * pr_img.cols] = (uc_pixel[0] / 255.0 - mean_value[0]) / std_value[0];
            uc_pixel += 3;
            ++i;
        }
    }

    // blob.convertTo(blob, CV_32FC3);
    // cv::cvtColor(blob, blob, cv::COLOR_BGR2RGB);
    // auto time_resize_end = std::chrono::high_resolution_clock::now();
    std::cout << "setBindingDimensions1:" << infer_h_ << "," << infer_w_ << std::endl;
    this->context_->setBindingDimensions(
        0,
        nvinfer1::Dims { 4,
            { 1, 3, infer_h_, infer_w_ } });
    std::cout << "setBindingDimensions2:" << infer_h_ << "," << infer_w_ << std::endl;
    cudaMemcpy(
        this->device_ptrs[0],
        data_,
        infer_w_ * infer_h_ * 3 * sizeof(float),
        cudaMemcpyHostToDevice);
    // auto time_cpy_end = std::chrono::high_resolution_clock::now();
    // auto time_resize = std::chrono::duration_cast<std::chrono::microseconds>(time_resize_end - time_cpy_start).count();
    // auto time_cpy = std::chrono::duration_cast<std::chrono::microseconds>(time_cpy_end - time_resize_end).count();
    // std::cout << "copy from mat resize time:" << time_resize << ", mem cpy:" << time_cpy << std::endl;
}

void LaneSeg::infer()
{
    auto time_infer_start = std::chrono::high_resolution_clock::now();
    this->context_->enqueueV2(
        this->device_ptrs.data(),
        this->stream_,
        nullptr);
    // auto time_equeue_end = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < this->num_outputs; i++) {
        size_t osize = this->output_bindings[i].size * this->output_bindings[i].dsize;
        cudaMemcpy(this->host_ptrs[i],
            this->device_ptrs[i + this->num_inputs],
            osize,
            cudaMemcpyDeviceToHost);
    }
    // auto time_cpy_end = std::chrono::high_resolution_clock::now();
    cudaStreamSynchronize(this->stream_);
    // auto time_sync_end = std::chrono::high_resolution_clock::now();
    // auto time_enqeue = std::chrono::duration_cast<std::chrono::microseconds>(time_equeue_end - time_infer_start).count();
    // auto time_cpy = std::chrono::duration_cast<std::chrono::microseconds>(time_cpy_end - time_equeue_end).count();
    // auto time_sync = std::chrono::duration_cast<std::chrono::microseconds>(time_sync_end - time_cpy_end).count();
    // std::cout << "infer enqeue:" << time_enqeue << ", mem cpy:" << time_cpy << ", sync:" << time_sync << std::endl;
}

void LaneSeg::get_mask(cv::Mat lane_mask)
{
    auto time_trt_start = std::chrono::high_resolution_clock::now();
    int32_t* net_out = static_cast<int32_t*>(this->host_ptrs[0]); // 1x1
    // int mask_h = infer_h_ / 8;
    // int mask_w = infer_w_ / 8;
    // int orig_image_h = image_orig_.rows;
    // int orig_image_w = image_orig_.cols;
    // cv::Mat score_map(infer_h_, infer_w_, CV_32FC());
    unsigned char* dst = lane_mask.ptr<unsigned char>();
    for (int i = 0; i < infer_h_; ++i) {
        // cv::Vec<unsigned char, LANE_NUM_CLASSES> *p1 = lane_mask.ptr<cv::Vec<unsigned char, LANE_NUM_CLASSES>>(i);
        for (int j = 0; j < infer_w_; ++j) {
            unsigned char max_class = 0;
            float max_score = -1;
            // for (int c = 0; c < num_class_; ++c)
            // {
            //     //  p1[j][c] =
            //     //     net_out[c * infer_w_ * infer_h_ + i * infer_w_ + j];
            //     float score = net_out[c * infer_w_ * infer_h_ + i * infer_w_ + j];
            //     if (max_score < score){
            //         max_score = score;
            //         max_class = c;
            //     }
            // }
            dst[i * infer_w_ + j] = net_out[i * infer_w_ + j];
        }
    }
}

void LaneSeg::process(const cv::Mat& image, std::vector<std::vector<cv::Point2f>>& contours)
{
    if (image.empty()) {
        std::cout << "lane seg image empty!!! " << std::endl;
        return;
    }

    float score_thres = 0.5; // 0.25f;
    {
        ++infer_count_;
        cv::Mat image_clone = image.clone();
        copy_from_mat(image_clone);
        infer();
        cv::Mat lane_mask;
        cv::Mat lane_masku(infer_h_, infer_w_, CV_8UC1);
        get_mask(lane_masku);
        lane_masku = lane_masku * 50;
        cv::resize(lane_masku, lane_masku, cv::Size(image.cols, image.rows));
        cv::imwrite("log/lane_mask_" + std::to_string(infer_count_) + "_orig.jpg", image_clone);
        std::vector<cv::Vec4i> hie;
        std::vector<std::vector<cv::Point>> cnts2d, cnts2d_orig;
        cv::Mat cls_mask;
        for (int icls = 1; icls < 5; ++icls) {
            hie.clear();
            cnts2d.clear();
            cls_mask = lane_masku == icls * 50;
            cv::findContours(cls_mask, cnts2d, hie, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
            for (auto& cnt : cnts2d) {
                std::vector<cv::Point2f> rcnt;
                std::vector<cv::Point> rcnt_orig;
                for (auto& pt : cnt) {
                    cv::Point2f ptf;
                    cv::Point ptd;
                    ptf.x = 1.0 * (pt.x - pad_left_) / (infer_w_ - 2 * pad_left_);
                    ptf.y = 1.0 * (pt.y - pad_top_) / (infer_h_ - 2 * pad_top_);
                    ptd.x = pt.x - pad_left_;
                    ptd.y = pt.y - pad_top_;
                    rcnt_orig.push_back(ptd);
                    if (ptf.x < 0 || ptf.x > 1 || ptf.y < 0 || ptf.y > 1) {
                        std::cout << "bad pt: " << ptf.x << "," << ptf.y << "; orig:" << pt.x << "," << pt.y << std::endl;
                    }
                    rcnt.push_back(ptf);
                }
                contours.push_back(std::move(rcnt));
                cnts2d_orig.push_back(rcnt_orig);
            }
        }
        cv::drawContours(image_clone, cnts2d_orig, -1, cv::Scalar { 50, 50, 50 }, 2, cv::LINE_8);
        cv::imwrite("log/lane_mask_" + std::to_string(infer_count_) + "_drawmask.jpg", image_clone);
        std::cout << "contour.size():" << contours.size() << std::endl;
        // cv::findContours(lane_mask, contours, cv::RETR_EXTERNAL, cv::CHAIN_APPROX_SIMPLE);
        // postprocess_cpu(objs, score_thres);
        // auto end_time = std::chrono::high_resolution_clock::now();
        // auto duration  = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);
        // time_total_ += duration.count();
        // std::cout << "total:" << time_total_ << ", count:" << infer_count_ << std::endl;
        // std::cout << "detect process time: " << 1.0 * time_total_ / infer_count_ << ", cur:" << duration.count()  << std::endl;
    }
}

std::shared_ptr<LaneSeg> LaneSeg::get_detector(const std::string& model_dir)
{
    if (!detector_) {
        detector_ = std::shared_ptr<LaneSeg>(new LaneSeg());
        detector_->init(model_dir);
    }
    return detector_;
}
